import numpy as np
import random

def euclidean_distance(x1, x2):
    """计算欧氏距离"""
    return np.sqrt(np.sum((x1 - x2) **2))

class KMeans:
    """k-means聚类算法实现"""
    
    def __init__(self, k=2, max_iters=100, tol=1e-4):
        self.k = k  # 聚类数量
        self.max_iters = max_iters  # 最大迭代次数
        self.tol = tol  # 收敛阈值
        self.centroids = None  # 质心
        self.clusters = None  # 聚类结果
    
    def fit(self, X):
        """
        训练k-means模型
        X: 输入数据，二维数组 (n_samples, n_features)
        """
        X = np.array(X)
        n_samples, n_features = X.shape
        
        # 随机初始化质心
        random_indices = random.sample(range(n_samples), self.k)
        self.centroids = [X[i] for i in random_indices]
        
        for _ in range(self.max_iters):
            # 初始化聚类
            self.clusters = [[] for _ in range(self.k)]
            
            # 将每个样本分配到最近的质心
            for idx, sample in enumerate(X):
                # 计算到每个质心的距离
                distances = [euclidean_distance(sample, centroid) for centroid in self.centroids]
                # 找到最近的质心索引
                closest_idx = np.argmin(distances)
                self.clusters[closest_idx].append(idx)
            
            # 保存当前质心
            prev_centroids = np.copy(self.centroids)
            
            # 更新质心（计算每个聚类的均值）
            for cluster_idx in range(self.k):
                cluster_samples = X[self.clusters[cluster_idx]]
                self.centroids[cluster_idx] = np.mean(cluster_samples, axis=0)
            
            # 检查是否收敛
            centroids_movement = [
                euclidean_distance(self.centroids[i], prev_centroids[i])
                for i in range(self.k)
            ]
            
            if np.max(centroids_movement) < self.tol:
                break  # 质心移动小于阈值，停止迭代
    
    def predict(self, X):
        """
        预测样本所属的聚类
        X: 输入数据，二维数组 (n_samples, n_features)
        返回：每个样本的聚类标签
        """
        X = np.array(X)
        predictions = []
        
        for sample in X:
            distances = [euclidean_distance(sample, centroid) for centroid in self.centroids]
            closest_idx = np.argmin(distances)
            predictions.append(closest_idx)
        
        return predictions

# 示例用法
if __name__ == "__main__":
    # 生成示例数据
    np.random.seed(42)
    # 生成三个聚类中心
    center1 = np.array([1, 1])
    center2 = np.array([5, 5])
    center3 = np.array([9, 1])
    
    # 围绕每个中心生成随机点
    cluster1 = np.random.randn(30, 2) + center1
    cluster2 = np.random.randn(30, 2) + center2
    cluster3 = np.random.randn(30, 2) + center3
    
    # 合并数据
    X = np.vstack((cluster1, cluster2, cluster3))
    
    # 创建并训练k-means模型
    kmeans = KMeans(k=3, max_iters=100)
    kmeans.fit(X)
    
    # 预测所有样本的聚类
    predictions = kmeans.predict(X)
    
    print("聚类质心:")
    for i, centroid in enumerate(kmeans.centroids):
        print(f"质心 {i}: {centroid}")
    
    print("\n前10个样本的聚类结果:", predictions[:10])
    
    # 计算每个聚类的样本数量
    cluster_counts = [len(cluster) for cluster in kmeans.clusters]
    print("\n每个聚类的样本数量:", cluster_counts)
